import os, time
from datetime import datetime
import streamlit as st
from pathlib import Path
from indexing import build_index, delete_index

# 配置常量
KNOWLEDGE_BASE_DIR = Path("knowledge_base")
ALLOWED_EXTENSIONS = [".txt", ".md", ".pdf", ".docx"]
MAX_FILE_SIZE_MB = 200
# 分类子目录
CATEGORIES = ["项目知识库", "业务知识库", "内部常见问题（FAQ）"]

def get_index_id(category):
    """根据分类获取索引ID"""
    if category == "项目知识库":
        return "project_docs"
    elif category == "业务知识库":
        return "business_rules"
    elif category == "内部常见问题（FAQ）":
        return "internal_faq"
    return None


def init_knowledge_base():
    """初始化知识库目录和session状态"""
    #  创建知识库目录
    KNOWLEDGE_BASE_DIR.mkdir(exist_ok=True) #只有在目录不存在时创建目录，目录已存在时不会抛出异常
    # 创建分类子目录
    for category in CATEGORIES:
        (KNOWLEDGE_BASE_DIR / category).mkdir(exist_ok=True)
    # 初始化session状态中的变量
    if "kb_files" not in st.session_state:
        st.session_state.kb_files = list_knowledge_base_files()
    if "show_upload_form" not in st.session_state:
        st.session_state.show_upload_form = False

def list_knowledge_base_files():
    """获取知识库文件列表"""
    files = []
    #使用 glob 方法遍历知识库目录下的所有文件和子目录
    for category in CATEGORIES:
        for f in (KNOWLEDGE_BASE_DIR / category).glob("*"):
            # 如果是文件且后缀名在允许的列表中
            if f.is_file() and f.suffix.lower() in ALLOWED_EXTENSIONS:
                # 获取文件的最后修改时间，并格式化为字符串
                mtime = datetime.fromtimestamp(f.stat().st_mtime)
                files.append({
                    "name": f.name,
                    "size": f.stat().st_size,
                    "create_time": mtime.strftime("%Y-%m-%d %H:%M:%S"),
                    "category": category,
                    "path": str(f)
                })

    # 返回排序后的文件列表，排序依据是文件的create_time，并且是降序排列
    return sorted(files, key=lambda x: x["create_time"], reverse=True)

def format_file_size(size_bytes):
    """格式化文件大小显示"""
    if size_bytes < 1024:
        return f"{size_bytes} B"
    elif size_bytes < 1024**2:
        return f"{size_bytes/1024:.1f} KB"
    else:
        return f"{size_bytes/(1024**2):.1f} MB"

def process_uploaded_files(uploaded_files, category):
    """处理上传的文件，将文件保存到知识库目录，并更新文件列表"""
    if not uploaded_files:
        st.toast("未选择任何文件", icon="⚠️")
        return
    for file in uploaded_files:
        # 检查文件大小
        if file.size > MAX_FILE_SIZE_MB * 1024 * 1024:
            st.error(f"文件 {file.name} 超过大小限制（{MAX_FILE_SIZE_MB}MB）")
            continue
        save_path = KNOWLEDGE_BASE_DIR / category / file.name
        try:
            with open(save_path, "wb") as f:
                f.write(file.getbuffer())
            # 创建索引
            with st.spinner("文件解析中..."):
                index_id = get_index_id(category)
                build_index(index_id, str(save_path))
            # 更新文件列表
            st.session_state.kb_files = list_knowledge_base_files()
            st.toast(f"成功上传：{file.name}", icon="✅")
        except Exception as e:
            st.error(f"上传失败：{str(e)}")
    # 完成文件上传后收起上传表单
    time.sleep(1)
    st.session_state.show_upload_form = False
    st.rerun()

def handle_file_upload():
    """处理文件上传逻辑"""
    if st.session_state.show_upload_form:
        # clear_on_submit=True，表示在表单提交后重置表单, 表单中的临时文件会被清空
        with st.form("upload_form", clear_on_submit=True):
            category = st.selectbox("选择分类", CATEGORIES, key="upload_category")
            uploaded_files = st.file_uploader(
                label="选择文件（支持多选）",
                type=[ext[1:] for ext in ALLOWED_EXTENSIONS],
                accept_multiple_files=True,
                help=f"最大单个文件大小：{MAX_FILE_SIZE_MB}MB"
            )
            submitted = st.form_submit_button("🚀 开始上传", type="primary")
            if submitted:
                # 点击上传按钮后，调用process_uploaded_files函数处理上传的文件
                process_uploaded_files(uploaded_files, category)

def display_file_operations():
    """分类筛选和搜索框展示"""
    col1, col2, col3 = st.columns([2, 7, 2])
    with col1:
        selected_category = st.selectbox("分类筛选", ["全部"] + CATEGORIES)
    with col2:
        search_term = st.text_input("🔍 搜索文件", placeholder="输入文件名...")
    with col3:
        st.markdown("")
        st.markdown("")
        button_name = "上传文件 ⏬"
        if st.session_state.show_upload_form:
            button_name = "收起表单 ⏫"
        if st.button(button_name, type="secondary" ):
            st.session_state.show_upload_form = not st.session_state.show_upload_form
            st.rerun()
    # 文件按搜索条件进行筛选
    filtered_files = [
        f for f in st.session_state.kb_files
        if (selected_category == "全部" or f["category"] == selected_category)
           and search_term.lower() in f["name"].lower()
    ]
    if not filtered_files:
        st.info("暂无知识库文件")
        return
    return filtered_files

def display_list(filtered_files):
    # 添加表头
    with st.container(border=True):
        col1, col2, col3, col4, col5, col6 = st.columns([3, 1, 1, 2, 1, 1])
        with col1:
            st.markdown("**文件名称**")
        with col2:
            st.markdown("**分类**")
        with col3:
            st.markdown("**大小**")
        with col4:
            st.markdown("**上传时间**")
        with col5:
            st.markdown("**下载**")  # 下载按钮
        with col6:
            st.markdown("**删除**")  # 删除按钮

    # 遍历文件列表，显示文件信息
    if filtered_files:
        for file in filtered_files:
            with st.container(border=True):
                col1, col2, col3, col4, col5, col6 = st.columns([3, 1, 1, 2, 1, 1])
                with col1:
                    st.markdown(f"**{file['name']}**    ")
                with col2:
                    st.markdown(f"{file['category']}    ")
                with col3:
                    st.markdown(f"{format_file_size(file['size'])}    ")
                with col4:
                    st.markdown(f"{file['create_time']}")
                with col5:
                    st.download_button(
                        label="⬇️ 下载",
                        data=open(file["path"], "rb").read(),
                        file_name=file["name"],
                        key=f"download_{file['name']}"
                    )
                with col6:
                    if st.button("❌ 删除", key=f"delete_{file['name']}"):
                        try:
                            delete_index(index_id = get_index_id(file["category"]), file_path=file["path"])
                            os.remove(file["path"])
                            # 更新文件列表
                            st.session_state.kb_files = list_knowledge_base_files()
                            st.rerun()
                        except Exception as e:
                            st.error(f"删除失败：{str(e)}")

def main():
    st.title("📁 知识库管理")
    init_knowledge_base()
    handle_file_upload()
    filter_files = display_file_operations()
    display_list(filter_files)
    
